From 7e80c684cad3249923b05ecbe02c009b111d9b9f Mon Sep 17 00:00:00 2001
From: Ewan Higgs <ewan_higgs@yahoo.co.uk>
Date: Mon, 27 Feb 2017 01:13:07 +0100
Subject: [PATCH] Fix for #3722

When using init, add [[bin]] and [lib] section where appropriate.
---
 src/cargo/ops/cargo_new.rs | 59 +++++++++++++++++++++++++++++---------
 tests/init.rs              | 53 +++++++++++++++++++++++++++-------
 2 files changed, 88 insertions(+), 24 deletions(-)

diff --git a/src/cargo/ops/cargo_new.rs b/src/cargo/ops/cargo_new.rs
index d211dc3e9..993e85b4d 100644
--- a/src/cargo/ops/cargo_new.rs
+++ b/src/cargo/ops/cargo_new.rs
@@ -10,6 +10,7 @@ use term::color::BLACK;
 
 use handlebars::{Handlebars, no_escape};
 use tempdir::TempDir;
+use toml;
 
 use core::Workspace;
 use sources::git::clone;
@@ -44,6 +45,7 @@ struct MkOptions<'a> {
     path: &'a Path,
     name: &'a str,
     bin: bool,
+    source_files: Vec<SourceFileInformation>,
 }
 
 impl Decodable for VersionControl {
@@ -127,11 +129,40 @@ fn get_input_template(config: &Config, opts: &MkOptions) -> CargoResult<Template
         // no template given, use either "lib" or "bin" templates depending on the
         // presence of the --bin flag.
         TemplateType::Builtin => {
-            let template_files = if opts.bin {
-                create_bin_template()
-            } else {
-                create_lib_template()
+
+
+            let mut cargotoml_path_specifier = String::new();
+            // If the main function is in a weird location, let Cargo.toml know.
+            // Otherwise, if the library files are somewhere else, let Cargo.toml know.
+            for i in &opts.source_files {
+                if i.bin {
+                    if i.relative_path != "src/main.rs" {
+                        cargotoml_path_specifier.push_str(&format!(r#"
+[[bin]]
+name = "{}"
+path = {}
+"#, i.target_name, toml::Value::String(i.relative_path.clone())));
+                    }
+                } else {
+                    if i.relative_path != "src/lib.rs" {
+                        cargotoml_path_specifier.push_str(&format!(r#"
+[lib]
+name = "{}"
+path = {}
+"#, i.target_name, toml::Value::String(i.relative_path.clone())));
+                    }
+                }
+            }
+
+            let mut template_files = create_generic_template(&cargotoml_path_specifier);
+            let mut source_files = match (opts.bin, cargotoml_path_specifier.len()) {
+                (true, 0)  => create_bin_template(),
+                (true, _)  => vec![],
+                (false, 0) => create_lib_template(),
+                (false, _)  => vec![],
             };
+            template_files.extend(source_files.drain(..));
+
             TemplateSet {
                 template_dir: None,
                 template_files: template_files
@@ -338,6 +369,7 @@ pub fn new(opts: NewOptions, config: &Config) -> CargoResult<()> {
         path: &path,
         name: name,
         bin: opts.bin,
+        source_files: vec![],
     };
 
     mk(config, &mkopts).chain_error(|| {
@@ -405,6 +437,7 @@ pub fn init(opts: NewOptions, config: &Config) -> CargoResult<()> {
         path: &path,
         name: name,
         bin: src_paths_types.iter().any(|x|x.bin),
+        source_files: src_paths_types,
     };
 
     mk(config, &mkopts).chain_error(|| {
@@ -659,40 +692,38 @@ fn walk_template_dir(dir: &Path, cb: &mut FnMut(DirEntry) -> CargoResult<()>) ->
 /// Create a generic template
 ///
 /// This consists of a Cargo.toml, and a src directory.
-fn create_generic_template() -> Vec<Box<TemplateFile>> {
-    let template_file = Box::new(InMemoryTemplateFile::new(PathBuf::from("Cargo.toml"),
-    String::from(r#"[package]
+fn create_generic_template(extra_cargo_info: &str) -> Vec<Box<TemplateFile>> {
+    let mut cargo_toml_contents = String::from(r#"[package]
 name = "{{name}}"
 version = "0.1.0"
 authors = [{{toml-escape author}}]
 
 [dependencies]
-"#)));
+"#);
+    cargo_toml_contents.push_str(extra_cargo_info);
+    let template_file = Box::new(InMemoryTemplateFile::new(PathBuf::from("Cargo.toml"),
+                            cargo_toml_contents));
     vec![template_file]
 }
 
 /// Create a new "lib" project
 fn create_lib_template() -> Vec<Box<TemplateFile>> {
-    let mut template_files = create_generic_template();
     let lib_file = Box::new(InMemoryTemplateFile::new(PathBuf::from("src/lib.rs"),
     String::from(r#"#[test]
 fn it_works() {
 }
 "#)));
-    template_files.push(lib_file);
-    template_files
+    vec![lib_file]
 }
 
 /// Create a new "bin" project
 fn create_bin_template() -> Vec<Box<TemplateFile>> {
-    let mut template_files = create_generic_template();
     let main_file = Box::new(InMemoryTemplateFile::new(PathBuf::from("src/main.rs"),
 String::from("fn main() {
     println!(\"Hello, world!\");
 }
 ")));
-    template_files.push(main_file);
-    template_files
+    vec![main_file]
 }
 
 #[cfg(test)]
diff --git a/tests/init.rs b/tests/init.rs
index c13d57d73..00d636be7 100644
--- a/tests/init.rs
+++ b/tests/init.rs
@@ -63,7 +63,7 @@ fn both_lib_and_bin() {
                     "[ERROR] can't specify both lib and binary outputs"));
 }
 
-fn bin_already_exists(explicit: bool, rellocation: &str) {
+fn bin_already_exists(explicit: bool, rellocation: &str, needs_bin_section: bool) {
     let path = paths::root().join("foo");
     fs::create_dir_all(&path.join("src")).unwrap();
 
@@ -94,36 +94,47 @@ fn bin_already_exists(explicit: bool, rellocation: &str) {
     let mut new_content = Vec::new();
     File::open(&sourcefile_path).unwrap().read_to_end(&mut new_content).unwrap();
     assert_eq!(Vec::from(content as &[u8]), new_content);
+
+	let mut cargo_content = String::new();
+    File::open(&paths::root().join("foo/Cargo.toml")).unwrap().read_to_string(&mut cargo_content).unwrap();
+	// Check that Cargo.toml has a bin section pointing to the correct location (if needed)
+	if needs_bin_section {
+		assert!(cargo_content.contains(r#"[[bin]]"#));
+        assert_that(&paths::root().join("foo/src/main.rs"), is_not(existing_file()));
+	} else {
+		assert!(!cargo_content.contains(r#"[[bin]]"#));
+        assert_that(&paths::root().join("foo/src/main.rs"), existing_file());
+	}
 }
 
 #[test]
 fn bin_already_exists_explicit() {
-    bin_already_exists(true, "src/main.rs")
+    bin_already_exists(true, "src/main.rs", false)
 }
 
 #[test]
 fn bin_already_exists_implicit() {
-    bin_already_exists(false, "src/main.rs")
+    bin_already_exists(false, "src/main.rs", false)
 }
 
 #[test]
 fn bin_already_exists_explicit_nosrc() {
-    bin_already_exists(true, "main.rs")
+    bin_already_exists(true, "main.rs", true)
 }
 
 #[test]
 fn bin_already_exists_implicit_nosrc() {
-    bin_already_exists(false, "main.rs")
+    bin_already_exists(false, "main.rs", true)
 }
 
 #[test]
 fn bin_already_exists_implicit_namenosrc() {
-    bin_already_exists(false, "foo.rs")
+    bin_already_exists(false, "foo.rs", true)
 }
 
 #[test]
 fn bin_already_exists_implicit_namesrc() {
-    bin_already_exists(false, "src/foo.rs")
+    bin_already_exists(false, "src/foo.rs", true)
 }
 
 #[test]
@@ -190,7 +201,7 @@ cannot automatically generate Cargo.toml as the main target would be ambiguous
     assert_that(&paths::root().join("foo/Cargo.toml"), is_not(existing_file()));
 }
 
-fn lib_already_exists(rellocation: &str) {
+fn lib_already_exists(rellocation: &str, needs_lib_section: bool) {
     let path = paths::root().join("foo");
     fs::create_dir_all(&path.join("src")).unwrap();
 
@@ -213,16 +224,38 @@ fn lib_already_exists(rellocation: &str) {
     let mut new_content = Vec::new();
     File::open(&sourcefile_path).unwrap().read_to_end(&mut new_content).unwrap();
     assert_eq!(Vec::from(content as &[u8]), new_content);
+
+	let mut cargo_content = String::new();
+    File::open(&paths::root().join("foo/Cargo.toml")).unwrap().read_to_string(&mut cargo_content).unwrap();
+	// Check that Cargo.toml has a lib section pointing to the correct location (if needed)
+	if needs_lib_section {
+		assert!(cargo_content.contains(r#"[lib]"#));
+        assert_that(&paths::root().join("foo/src/lib.rs"), is_not(existing_file()));
+	} else {
+		assert!(!cargo_content.contains(r#"[lib]"#));
+        assert_that(&paths::root().join("foo/src/lib.rs"), existing_file());
+	}
+
 }
 
 #[test]
 fn lib_already_exists_src() {
-    lib_already_exists("src/lib.rs")
+    lib_already_exists("src/lib.rs", false)
 }
 
 #[test]
 fn lib_already_exists_nosrc() {
-    lib_already_exists("lib.rs")
+    lib_already_exists("lib.rs", true)
+}
+
+#[test]
+fn no_lib_already_exists_src_add_lib_section() {
+    lib_already_exists("src/foo.rs", true)
+}
+
+#[test]
+fn no_lib_already_exists_nosrc_add_lib_section() {
+    lib_already_exists("foo.rs", true)
 }
 
 #[test]
-- 
2.30.2